{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import dataclasses\n",
    "import mediapy\n",
    "from huggingface_hub import PyTorchModelHubMixin\n",
    "from huggingface_hub import ModelCard\n",
    "from gpudrive.networks.late_fusion import NeuralNet\n",
    "\n",
    "from gpudrive.env.config import EnvConfig\n",
    "from gpudrive.env.env_torch import GPUDriveTorchEnv\n",
    "from gpudrive.visualize.utils import img_from_fig\n",
    "from gpudrive.env.dataset import SceneDataLoader\n",
    "from gpudrive.utils.config import load_config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'max_controlled_agents': 64, 'ego_state': True, 'road_map_obs': True, 'partner_obs': True, 'norm_obs': True, 'remove_non_vehicles': True, 'lidar_obs': False, 'reward_type': 'weighted_combination', 'collision_weight': -0.75, 'off_road_weight': -0.75, 'goal_achieved_weight': 1.0, 'dynamics_model': 'classic', 'collision_behavior': 'ignore', 'dist_to_goal_threshold': 2.0, 'polyline_reduction_threshold': 0.1, 'sampling_seed': 42, 'obs_radius': 50.0, 'action_space_steer_disc': 13, 'action_space_accel_disc': 7, 'init_mode': 'all_non_trivial'}\n"
     ]
    }
   ],
   "source": [
    "# Configs model has been trained with\n",
    "config = load_config(\"../../examples/experimental/config/reliable_agents_params\")\n",
    "\n",
    "print(config)\n",
    "\n",
    "max_agents = config.max_controlled_agents\n",
    "num_envs = 2\n",
    "device = \"cpu\" # cpu just because we're in a notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load pre-trained agent via Hugging Face hub\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_agent = NeuralNet.from_pretrained(\"daphne-cornelisse/policy_S10_000_02_27\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "91"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Agent has an action dimension of 91: 13 steering wheel angle discretizations x 9 acceleration discretizations\n",
    "sim_agent.action_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2984"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Size of flattened observation vector\n",
    "sim_agent.obs_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ffn', 'model_hub_mixin', 'pytorch_model_hub_mixin']"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Some other info\n",
    "card = ModelCard.load(\"daphne-cornelisse/policy_S10_000_02_27\")\n",
    "card.data.tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model architecture\n",
    "#agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weights \n",
    "#agent.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create data loader\n",
    "train_loader = SceneDataLoader(\n",
    "    root='../../data/processed/examples',\n",
    "    batch_size=num_envs,\n",
    "    dataset_size=100,\n",
    "    sample_with_replacement=False,\n",
    ")\n",
    "\n",
    "# Set params\n",
    "env_config = dataclasses.replace(\n",
    "    EnvConfig(),\n",
    "    ego_state=config.ego_state,\n",
    "    road_map_obs=config.road_map_obs,\n",
    "    partner_obs=config.partner_obs,\n",
    "    reward_type=config.reward_type,\n",
    "    norm_obs=config.norm_obs,\n",
    "    dynamics_model=config.dynamics_model,\n",
    "    collision_behavior=config.collision_behavior,\n",
    "    dist_to_goal_threshold=config.dist_to_goal_threshold,\n",
    "    polyline_reduction_threshold=config.polyline_reduction_threshold,\n",
    "    remove_non_vehicles=config.remove_non_vehicles,\n",
    "    lidar_obs=config.lidar_obs,\n",
    "    disable_classic_obs=config.lidar_obs,\n",
    "    obs_radius=config.obs_radius,\n",
    "    steer_actions = torch.round(\n",
    "        torch.linspace(-torch.pi, torch.pi, config.action_space_steer_disc), decimals=3  \n",
    "    ),\n",
    "    accel_actions = torch.round(\n",
    "        torch.linspace(-4.0, 4.0, config.action_space_accel_disc), decimals=3\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Make env\n",
    "env = GPUDriveTorchEnv(\n",
    "    config=env_config,\n",
    "    data_loader=train_loader,\n",
    "    max_cont_agents=config.max_controlled_agents,\n",
    "    device=device,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['../../data/processed/examples/tfrecord-00000-of-01000_325.json',\n",
       " '../../data/processed/examples/tfrecord-00000-of-01000_4.json']"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.data_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 64, 2984])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next_obs = env.reset()\n",
    "\n",
    "control_mask = env.cont_agent_mask\n",
    "\n",
    "next_obs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "action, logprob, entropy, value = sim_agent(\n",
    "    next_obs[control_mask], deterministic=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([13]), torch.Size([13]), torch.Size([13]), torch.Size([13, 1]))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "action.shape, logprob.shape, entropy.shape, value.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rollout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 64, 2984])\n",
      "Step: 54"
     ]
    }
   ],
   "source": [
    "next_obs = env.reset()\n",
    "\n",
    "control_mask = env.cont_agent_mask\n",
    "\n",
    "print(next_obs.shape)\n",
    "\n",
    "frames = {f\"env_{i}\": [] for i in range(num_envs)}\n",
    "\n",
    "for time_step in range(env.episode_len):\n",
    "    print(f\"\\rStep: {time_step}\", end=\"\", flush=True)\n",
    "\n",
    "    # Predict actions\n",
    "    action, _, _, _ = sim_agent(\n",
    "        next_obs[control_mask], deterministic=False\n",
    "    )\n",
    "    action_template = torch.zeros(\n",
    "        (num_envs, max_agents), dtype=torch.int64, device=device\n",
    "    )\n",
    "    action_template[control_mask] = action.to(device)\n",
    "\n",
    "    # Step\n",
    "    env.step_dynamics(action_template)\n",
    "\n",
    "    # Render    \n",
    "    sim_states = env.vis.plot_simulator_state(\n",
    "        env_indices=list(range(num_envs)),\n",
    "        time_steps=[time_step]*num_envs,\n",
    "        zoom_radius=70,\n",
    "    )\n",
    "    \n",
    "    for i in range(num_envs):\n",
    "        frames[f\"env_{i}\"].append(img_from_fig(sim_states[i])) \n",
    "\n",
    "    next_obs = env.get_obs()\n",
    "    reward = env.get_rewards()\n",
    "    done = env.get_dones()\n",
    "    info = env.get_infos()\n",
    "    \n",
    "    if done.all():\n",
    "        break\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"show_videos\" style=\"border-spacing:0px;\"><tr><td style=\"padding:1px;\"><div style=\"display:flex; align-items:left;\">\n",
       "      <div style=\"display:flex; flex-direction:column; align-items:center;\">\n",
       "      <div>env_0</div><div><img width=\"500\" height=\"500\" style=\"image-rendering:auto; object-fit:cover;\" src=\"\"/></div></div></div></td><td style=\"padding:1px;\"><div style=\"display:flex; align-items:left;\">\n",
       "      <div style=\"display:flex; flex-direction:column; align-items:center;\">\n",
       "      <div>env_1</div><div><img width=\"500\" height=\"500\" style=\"image-rendering:auto; object-fit:cover;\" src=\"\"/></div></div></div></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mediapy.show_videos(frames, fps=15, width=500, height=500, columns=2, codec='gif')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gpudrive",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
